"""
RagFlow 服务模块
提供与 RagFlow 相关的业务逻辑
"""

import os
from typing import List, Optional, Tuple, Dict, Any, Union

from ragflow_client.api import RagFlowAPI
from ragflow_client.config import RagFlowConfig
from ragflow_client.utils.file_util import get_files, split_files_into_batches
from ragflow_client.utils.logger import get_logger

logger = get_logger(__name__)


class RagFlowService:
    """RagFlow 服务类"""
    
    def __init__(self, config: Optional[RagFlowConfig] = None):
        """
        初始化 RagFlow 服务
        
        Args:
            config: RagFlow 配置对象，如果为 None 则从环境变量加载
        """
        self.config = config or RagFlowConfig.from_env()
        self.api = RagFlowAPI(self.config)
    
    def upload_files_batch(self, batch_size: int = 10) -> Tuple[bool, str]:
        """
        批量上传文件到知识库
        
        Args:
            batch_size: 每批次上传的文件数量
            
        Returns:
            (成功状态, 消息) 元组
        """
        logger.info("开始批量上传文件到知识库")
        
        # 检查 API 连接
        status, msg = self.api.check_connection()
        if not status:
            logger.error(f"API 连接检查失败: {msg}")
            return False, msg
        
        # 获取符合条件的文件
        try:
            files = get_files(self.config.doc_dir, self.config.doc_suffix)
        except Exception as e:
            error_msg = f"获取文件失败: {str(e)}"
            logger.error(error_msg)
            return False, error_msg
        
        logger.info(f"找到 {len(files)} 个文件")
        
        if not files:
            logger.warning(f"在目录 {self.config.doc_dir} 中没有找到符合条件的文件")
            return True, "没有找到符合条件的文件"
        
        # 分批上传文件
        batches = split_files_into_batches(files, batch_size)
        
        success_count = 0
        failed_count = 0
        
        for i, batch in enumerate(batches):
            logger.info(f"上传批次 {i+1}/{len(batches)}, 文件数: {len(batch)}")
            
            status, result = self.api.upload_files(batch)
            
            if status:
                success_count += len(batch)
                logger.info(f"批次 {i+1} 上传成功")
            else:
                failed_count += len(batch)
                logger.error(f"批次 {i+1} 上传失败: {result}")
        
        logger.info(f"文件上传完成，成功: {success_count}，失败: {failed_count}")
        
        if failed_count > 0:
            return False, f"部分文件上传失败，成功: {success_count}，失败: {failed_count}"
        
        return True, f"所有文件上传成功，共 {success_count} 个文件"
    
    # ==================== 知识库管理 ====================
    
    def list_datasets(self, page: int = 1, page_size: int = 10) -> Tuple[bool, Union[List[Dict[str, Any]], str]]:
        """
        获取知识库列表
        
        Args:
            page: 页码，从 1 开始
            page_size: 每页数量
            
        Returns:
            (成功状态, 知识库列表或错误消息) 元组
        """
        logger.info("获取知识库列表")
        
        status, result = self.api.list_datasets(page, page_size)
        
        if status:
            # 处理不同的数据结构情况
            if isinstance(result, dict):
                # 尝试获取数据，处理可能的不同结构
                if "data" in result:
                    data = result["data"]
                    if isinstance(data, dict) and "items" in data:
                        datasets = data["items"]
                    elif isinstance(data, list):
                        datasets = data
                    else:
                        datasets = []
                else:
                    datasets = []
            elif isinstance(result, list):
                # 直接使用列表作为数据集
                datasets = result
            else:
                datasets = []
                
            logger.info(f"获取到 {len(datasets)} 个知识库")
            return True, datasets
        else:
            logger.error(f"获取知识库列表失败: {result}")
            return False, result
    
    def get_dataset_info(self, dataset_id: str = None) -> Tuple[bool, Union[Dict[str, Any], str]]:
        """
        获取知识库详情
        
        Args:
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            
        Returns:
            (成功状态, 知识库详情或错误消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"获取知识库详情: {dataset_id}")
        
        status, result = self.api.get_dataset(dataset_id)
        
        if status:
            dataset_info = result.get("data", {})
            logger.info(f"获取到知识库详情: {dataset_info.get('name')}")
            return True, dataset_info
        else:
            logger.error(f"获取知识库详情失败: {result}")
            return False, result
    
    def create_dataset(self, name: str, description: str = "", embedding_model: str = "text2vec-base") -> Tuple[bool, Union[Dict[str, Any], str]]:
        """
        创建知识库
        
        Args:
            name: 知识库名称
            description: 知识库描述
            embedding_model: 嵌入模型名称
            
        Returns:
            (成功状态, 知识库信息或错误消息) 元组
        """
        logger.info(f"创建知识库: {name}")
        
        status, result = self.api.create_dataset(name, description, embedding_model)
        
        if status:
            dataset_info = result.get("data", {})
            logger.info(f"知识库创建成功: {dataset_info.get('id')}")
            return True, dataset_info
        else:
            logger.error(f"创建知识库失败: {result}")
            return False, result
    
    def update_dataset(self, dataset_id: str = None, name: str = None, description: str = None) -> Tuple[bool, Union[Dict[str, Any], str]]:
        """
        更新知识库
        
        Args:
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            name: 知识库名称
            description: 知识库描述
            
        Returns:
            (成功状态, 知识库信息或错误消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"更新知识库: {dataset_id}")
        
        status, result = self.api.update_dataset(dataset_id, name, description)
        
        if status:
            dataset_info = result.get("data", {})
            logger.info(f"知识库更新成功: {dataset_info.get('name')}")
            return True, dataset_info
        else:
            logger.error(f"更新知识库失败: {result}")
            return False, result
    
    def delete_dataset(self, dataset_id: str = None) -> Tuple[bool, str]:
        """
        删除知识库
        
        Args:
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            
        Returns:
            (成功状态, 消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"删除知识库: {dataset_id}")
        
        status, result = self.api.delete_dataset(dataset_id)
        
        if status:
            logger.info(f"知识库删除成功: {dataset_id}")
            return True, f"知识库删除成功: {dataset_id}"
        else:
            logger.error(f"删除知识库失败: {result}")
            return False, result
    
    # ==================== 文档解析 ====================
    
    def parse_document(self, document_id: str, dataset_id: str = None, parser_id: str = None) -> Tuple[bool, Union[Dict[str, Any], str]]:
        """
        解析文档，将文档转换为知识块
        
        Args:
            document_id: 文档 ID
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            parser_id: 解析器 ID，如果为 None 则使用配置中的解析器 ID
            
        Returns:
            (成功状态, 解析结果或错误消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        parser_id = parser_id or self.config.parser_id
        logger.info(f"解析文档: {document_id}")
        
        status, result = self.api.parse_document(dataset_id, document_id, parser_id)
        
        if status:
            parse_result = result.get("data", {})
            logger.info(f"文档解析成功: {document_id}")
            return True, parse_result
        else:
            logger.error(f"解析文档失败: {result}")
            return False, result
    
    def list_chunks(self, dataset_id: str = None, document_id: str = None, page: int = 1, page_size: int = 10) -> Tuple[bool, Union[List[Dict[str, Any]], str]]:
        """
        获取知识块列表
        
        Args:
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            document_id: 文档 ID，如果指定则只返回该文档的知识块
            page: 页码，从 1 开始
            page_size: 每页数量
            
        Returns:
            (成功状态, 知识块列表或错误消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"获取知识库 {dataset_id} 的知识块列表")
        
        status, result = self.api.list_chunks(dataset_id, document_id, page, page_size)
        
        if status:
            chunks = result.get("data", {}).get("items", [])
            logger.info(f"获取到 {len(chunks)} 个知识块")
            return True, chunks
        else:
            logger.error(f"获取知识块列表失败: {result}")
            return False, result
    
    def get_chunk_info(self, chunk_id: str, dataset_id: str = None) -> Tuple[bool, Union[Dict[str, Any], str]]:
        """
        获取知识块详情
        
        Args:
            chunk_id: 知识块 ID
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            
        Returns:
            (成功状态, 知识块详情或错误消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"获取知识块详情: {chunk_id}")
        
        status, result = self.api.get_chunk(chunk_id, dataset_id)
        
        if status:
            chunk_info = result.get("data", {})
            logger.info(f"获取到知识块详情: {chunk_id}")
            return True, chunk_info
        else:
            logger.error(f"获取知识块详情失败: {result}")
            return False, result
    
    def delete_chunk(self, chunk_id: str, dataset_id: str = None) -> Tuple[bool, str]:
        """
        删除知识块
        
        Args:
            chunk_id: 知识块 ID
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            
        Returns:
            (成功状态, 消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"删除知识块: {chunk_id}")
        
        status, result = self.api.delete_chunk(chunk_id, dataset_id)
        
        if status:
            logger.info(f"知识块删除成功: {chunk_id}")
            return True, f"知识块删除成功: {chunk_id}"
        else:
            logger.error(f"删除知识块失败: {result}")
            return False, result
    
    def create_chunk(self, content: str, dataset_id: str = None, metadata: Dict[str, Any] = None) -> Tuple[bool, Union[Dict[str, Any], str]]:
        """
        创建知识块
        
        Args:
            content: 知识块内容
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            metadata: 知识块元数据
            
        Returns:
            (成功状态, 知识块信息或错误消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"创建知识块")
        
        status, result = self.api.create_chunk(dataset_id, content, metadata)
        
        if status:
            chunk_info = result.get("data", {})
            logger.info(f"知识块创建成功: {chunk_info.get('id')}")
            return True, chunk_info
        else:
            logger.error(f"创建知识块失败: {result}")
            return False, result
    
    def update_chunk(self, chunk_id: str, content: str = None, dataset_id: str = None, metadata: Dict[str, Any] = None) -> Tuple[bool, Union[Dict[str, Any], str]]:
        """
        更新知识块
        
        Args:
            chunk_id: 知识块 ID
            content: 知识块内容
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            metadata: 知识块元数据
            
        Returns:
            (成功状态, 知识块信息或错误消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"更新知识块: {chunk_id}")
        
        status, result = self.api.update_chunk(chunk_id, dataset_id, content, metadata)
        
        if status:
            chunk_info = result.get("data", {})
            logger.info(f"知识块更新成功: {chunk_id}")
            return True, chunk_info
        else:
            logger.error(f"更新知识块失败: {result}")
            return False, result
    
    # ==================== 文档管理 ====================
    
    def list_documents(self, dataset_id: str = None, page: int = 1, page_size: int = 10) -> Tuple[bool, Union[List[Dict[str, Any]], str]]:
        """
        获取文档列表
        
        Args:
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            page: 页码，从 1 开始
            page_size: 每页数量
            
        Returns:
            (成功状态, 文档列表或错误消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"获取知识库 {dataset_id} 的文档列表")
        
        status, result = self.api.list_documents(dataset_id, page, page_size)
        
        if status:
            documents = result.get("data", {}).get("items", [])
            logger.info(f"获取到 {len(documents)} 个文档")
            return True, documents
        else:
            logger.error(f"获取文档列表失败: {result}")
            return False, result
    
    def get_document_info(self, document_id: str, dataset_id: str = None) -> Tuple[bool, Union[Dict[str, Any], str]]:
        """
        获取文档详情
        
        Args:
            document_id: 文档 ID
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            
        Returns:
            (成功状态, 文档详情或错误消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"获取文档详情: {document_id}")
        
        status, result = self.api.get_document(document_id, dataset_id)
        
        if status:
            document_info = result.get("data", {})
            logger.info(f"获取到文档详情: {document_info.get('name')}")
            return True, document_info
        else:
            logger.error(f"获取文档详情失败: {result}")
            return False, result
    
    def delete_document(self, document_id: str, dataset_id: str = None) -> Tuple[bool, str]:
        """
        删除文档
        
        Args:
            document_id: 文档 ID
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            
        Returns:
            (成功状态, 消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"删除文档: {document_id}")
        
        status, result = self.api.delete_document(document_id, dataset_id)
        
        if status:
            logger.info(f"文档删除成功: {document_id}")
            return True, f"文档删除成功: {document_id}"
        else:
            logger.error(f"删除文档失败: {result}")
            return False, result
    
    # ==================== 检索与对话 ====================
    
    def search(self, query: str, dataset_id: str = None, top_k: int = 5, filter_threshold: float = 0.0) -> Tuple[bool, Union[List[Dict[str, Any]], str]]:
        """
        搜索知识库
        
        Args:
            query: 查询文本
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            top_k: 返回结果数量
            filter_threshold: 过滤阈值，相似度低于该值的结果将被过滤
            
        Returns:
            (成功状态, 搜索结果或错误消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"搜索知识库 {dataset_id}: {query}")
        
        status, result = self.api.search(query, dataset_id, top_k, filter_threshold)
        
        if status:
            search_results = result.get("data", [])
            logger.info(f"获取到 {len(search_results)} 条搜索结果")
            return True, search_results
        else:
            logger.error(f"搜索失败: {result}")
            return False, result
    
    def chat(self, query: str, dataset_id: str = None, history: List[Dict[str, str]] = None, 
             stream: bool = False, top_k: int = 5, filter_threshold: float = 0.0) -> Tuple[bool, Union[Dict[str, Any], str]]:
        """
        与知识库进行对话
        
        Args:
            query: 用户问题
            dataset_id: 知识库 ID，如果为 None 则使用配置中的知识库 ID
            history: 对话历史，格式为 [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
            stream: 是否使用流式响应
            top_k: 检索结果数量
            filter_threshold: 过滤阈值，相似度低于该值的结果将被过滤
            
        Returns:
            (成功状态, 对话结果或错误消息) 元组
        """
        dataset_id = dataset_id or self.config.kb_id
        logger.info(f"与知识库 {dataset_id} 对话: {query}")
        
        status, result = self.api.chat(query, dataset_id, history, stream, top_k, filter_threshold)
        
        if status:
            if stream:
                logger.info("获取到流式对话响应")
                return True, result
            else:
                chat_result = result.get("data", {})
                logger.info(f"获取到对话响应: {chat_result.get('answer', '')[:50]}...")
                return True, chat_result
        else:
            logger.error(f"对话失败: {result}")
            return False, result
    
    # ==================== 系统信息 ====================
    
    def get_parsers(self) -> Tuple[bool, Union[List[Dict[str, Any]], str]]:
        """
        获取解析器列表
        
        Returns:
            (成功状态, 解析器列表或错误消息) 元组
        """
        logger.info("获取解析器列表")
        
        status, result = self.api.list_parsers()
        
        if status:
            parsers = result.get("data", [])
            logger.info(f"获取到 {len(parsers)} 个解析器")
            return True, parsers
        else:
            logger.error(f"获取解析器列表失败: {result}")
            return False, result
    
    def get_embedding_models(self) -> Tuple[bool, Union[List[Dict[str, Any]], str]]:
        """
        获取嵌入模型列表
        
        Returns:
            (成功状态, 嵌入模型列表或错误消息) 元组
        """
        logger.info("获取嵌入模型列表")
        
        status, result = self.api.list_embedding_models()
        
        if status:
            models = result.get("data", [])
            logger.info(f"获取到 {len(models)} 个嵌入模型")
            return True, models
        else:
            logger.error(f"获取嵌入模型列表失败: {result}")
            return False, result